package com.ckw.question.utils;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Configuration;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;
import redis.clients.jedis.exceptions.JedisException;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;

@Configuration
public class RedisRateLimiter {

    private static final String LUA_SCRIPT_PATH = "token_bucket.lua";
    private static final String REDIS_HOST = "192.168.116.135";
    private static final int REDIS_PORT = 6379;
    private static final String REDIS_PASSWORD = "root";
    private static final String TOKEN_BUCKET_KEY_PREFIX = "token_bucket:";
    private static final String TOKENS_KEY_SUFFIX = ":tokens";
    private static final String LAST_REFILL_TIME_KEY_SUFFIX = ":last_refill_time";

    private static JedisPool jedisPool; // Redis 连接池
    private static String scriptSha;    // Lua 脚本的 SHA1 哈希值

    // 静态代码块：在类加载时初始化连接池和 Lua 脚本
    static {
        initJedisPool();
        preloadLuaScript();
    }

    /**
     * 初始化 Redis 连接池
     */
    private static void initJedisPool() {
        JedisPoolConfig poolConfig = new JedisPoolConfig();
        jedisPool = new JedisPool(poolConfig, REDIS_HOST, REDIS_PORT, 2000, REDIS_PASSWORD);
        System.out.println("Redis 连接池初始化成功");
    }

    /**
     * 预加载 Lua 脚本并保存 SHA1 哈希值
     */
    private static void preloadLuaScript() {
        try (Jedis jedis = jedisPool.getResource()) {
            String luaScript = readLuaScript();
            scriptSha = jedis.scriptLoad(luaScript); // 加载脚本并获取 SHA1
            System.out.println("Lua 脚本加载成功，SHA1: " + scriptSha);
        } catch (JedisException | IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 读取 Lua 脚本文件
     */
    private static String readLuaScript() throws IOException {
        StringBuilder sb = new StringBuilder();
        try (InputStream inputStream = RedisRateLimiter.class.getClassLoader().getResourceAsStream(LUA_SCRIPT_PATH);
             BufferedReader br = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            while ((line = br.readLine()) != null) {
                sb.append(line).append("\n");
            }
        }
        return sb.toString();
    }

    /**
     * 尝试消耗令牌
     *
     * tokensKey         令牌数量的键名
     * lastRefillTimeKey 上次补充时间的键名
     * @param capacity          令牌桶容量
     * @param refillRate        补充速率（令牌/秒）
     * @param tokensToConsume   尝试消耗的令牌数量
     * @return 是否允许请求
     */
    public boolean tryConsumeTokens(int capacity, double refillRate, int tokensToConsume) {
        try (Jedis jedis = jedisPool.getResource()) {
            // 准备脚本参数
            String tokensKey = TOKEN_BUCKET_KEY_PREFIX + TOKENS_KEY_SUFFIX;
            String lastRefillTimeKey = TOKEN_BUCKET_KEY_PREFIX + LAST_REFILL_TIME_KEY_SUFFIX;
            String[] keys = {tokensKey, lastRefillTimeKey};
            String[] argv = {String.valueOf(capacity), String.valueOf(refillRate), String.valueOf(tokensToConsume)};

            // 使用 SHA1 执行 Lua 脚本
            try {
                Object result = jedis.evalsha(scriptSha, keys.length, keys[0], keys[1], argv[0], argv[1], argv[2]);
                return result instanceof Long && (Long) result == 1;
            } catch (JedisException e) {
                if (e.getMessage().contains("NOSCRIPT")) {
                    // 如果脚本不存在，重新加载脚本
                    preloadLuaScript();
                    // 重试
                    Object result = jedis.evalsha(scriptSha, keys.length, keys[0], keys[1], argv[0], argv[1], argv[2]);
                    return result instanceof Long && (Long) result == 1;
                } else {
                    throw e;
                }
            }
        } catch (JedisException e) {
            e.printStackTrace();
            return false;
        }
    }

    /**
     * 关闭连接池
     */
    public static void close() {
        if (jedisPool != null) {
            jedisPool.close();
            System.out.println("Redis 连接池已关闭");
        }
    }
}